from ..model import build_detector
from ..dataset import build_dataset, build_preprocess_op
from torch.utils.data.dataloader import DataLoader
from progressbar import *

import torch
import os
import numpy as np



def cal_iou(bboxes1, bboxes2):
    max_x1 = np.maximum(bboxes1[:, 0], bboxes2[:, 0])
    max_y1 = np.maximum(bboxes1[:, 1], bboxes2[:, 1])
    min_x2 = np.minimum(bboxes1[:, 2], bboxes2[:, 2])
    min_y2 = np.minimum(bboxes1[:, 3], bboxes2[:, 3])

    w = np.maximum(min_x2 - max_x1, 0)
    h = np.maximum(min_y2 - max_y1, 0)
    inter_area = np.float(w * h)

    area1 = np.float((bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1]))
    area2 = np.float((bboxes2[:, 2] - bboxes2[:, 0]) * (bboxes2[:, 3] - bboxes2[:, 1]))

    iou = inter_area / (area1 + area2 - inter_area)
    return iou


def test(cfg):
    model_cfg = cfg['Model']
    dataset_cfg = cfg['Dataset']
    test_cfg = cfg['Test_cfg']

    check_point_path = test_cfg['check_point']
    model_cfg['mode'] = 'test'
    model = build_detector(model_cfg).cuda()
    model.load_state_dict(torch.load(check_point_path))
    model.eval()

    dataset_cfg['mode'] = 'test'
    dataset = build_dataset(dataset_cfg)

    data_pipe_cfgs = test_cfg['data_pipe']
    batch_size = test_cfg['batch_size']

    data_pipe = []
    for data_preprocess_cfg in data_pipe_cfgs:
        data_pipe.append(build_preprocess_op(data_preprocess_cfg))
    dataset.set_data_pipe(data_pipe)

    data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False, drop_last=False)

    sample_count = 0
    cls_true_count = 0
    loc_true_count = 0

    widgets = ['testing: ', Percentage(), ' ', Bar('|'), ' ', Timer(),
               ' ', ETA(), ' ']
    pbar = ProgressBar(widgets=widgets, maxval=len(data_loader)).start()
    for i, data in enumerate(data_loader):
        input_data, annotation = data
        input_data = input_data.cuda()
        cls_label = annotation[:, 0].numpy()
        bbox_label = annotation[:, 1:].numpy()
        clses, bboxes = model(input_data)
        clses = clses.cpu().detach().numpy()
        clses = np.argmax(clses, axis=1)
        bboxes = bboxes.cpu().detach().numpy()

        sample_count += cls_label.shape[0]
        cls_true_indexes = (cls_label - clses == 0)
        cls_true_count += np.sum(cls_true_indexes)

        ious = cal_iou(bbox_label, bboxes)
        loc_true_indexes = (ious >= 0.5)
        # loc_true_indexes = loc_true_indexes & cls_true_indexes
        loc_true_count += np.sum(loc_true_indexes)
        pbar.update(sample_count)
    pbar.finish()
    print('cls_rate:{} || loc_rate:{}'.format(cls_true_count/sample_count, loc_true_count/sample_count))


